import re
import numpy as np
import time

class MD_moment:
    def __init__(self, t, box_xyz, V, atom_num, table_headers, header_lines, atom_data):
        self.t = t  # Could represent time or temperature, based on your context
        self.box_xyz = box_xyz  # This represent dimensions of a box in 3D space
        self.V = V  # Possibly volume or velocity
        self.atom_num=atom_num
        self.table_headers=table_headers
        self.header_lines = header_lines  # General information about a system
        self.atom_data = atom_data  # atom_data is a np.array format

    ## method
    
    def sort(self, sort_column):
        if isinstance(sort_column, int):
            # If sort_column is an integer, we need to transform it to the field name
            sort_column = self.table_headers[sort_column]
        # Perform the sort using the field name directly
        self.atom_data = np.sort(self.atom_data, order=[sort_column])
    def col_search(self,col_name):
        if col_name in self.table_headers:
            return self.table_headers.index(col_name) 
    def add(self, col_name, col_list):
        # check if col_name in self.table_headers
        if col_name not in self.table_headers:
            from numpy.lib.recfunctions import append_fields
            self.table_headers.append(col_name)  # Update the headers list
            # Add a new column to the structured array
            # new_dtype = [(col_name, type(col_list[0]))]  # Determine the type based on the first element
            new_col_array = np.array(col_list, dtype=type(col_list[0]))
            self.atom_data = append_fields(self.atom_data, col_name, new_col_array, usemask=False)
        else:
            print(f'{col_name} in the tabel already.')
    def delete(self, col_name):
        from numpy.lib.recfunctions import drop_fields
        if col_name in self.table_headers:
            self.atom_data = drop_fields(self.atom_data, col_name, usemask=False)
            self.table_headers.remove(col_name)  # Also update the headers list
        else:
            print("Column not found") 
    def mol_add(self, chain_num):
        # if mol list already exists, delete it.
        if 'mol' in self.table_headers:
            self.delete('mol')
        else:
            # 分组并填充每个块的序列数
            mol_list = [0] * self.atom_num  # 初始化存储每个原子链段编号的列表
            atoms_per_chain = self.atom_num // chain_num
            for i in range(chain_num):
                start_index = i * atoms_per_chain
                end_index = start_index + atoms_per_chain
                if i == chain_num - 1:  # 最后一个block可能包含更多的元素
                    end_index = self.atom_num
                for j in range(start_index, end_index):
                    mol_list[j] = i + 1
            self.add('mol',mol_list)
    
    ## 两个序参量函数
    
    def orderPara_add(self, chain_num=0, R_search=7):

        class Point_Vector:
            def __init__(self, index, loca, vec, unit):
                self.index = index
                self.loca = loca
                self.vec = vec
                self.unit = unit
                self.ord = .0
                self.PB_cross_flag=[False, False, False]

        def orderCalcu(atom_index, loca_vec,dirc_vec,R_search,box):
            def cos(vector1, vector2):
                # Calculate the norms (lengths) of the vectors  
                norm_vector1 = np.linalg.norm(vector1)
                norm_vector2 = np.linalg.norm(vector2)
                # Calculate the cosine similarity
                return np.dot(vector1, vector2) / (norm_vector1 * norm_vector2) if norm_vector1 * norm_vector2!=0 else 0
            
            box_min_vector = [box[0][0], box[1][0], box[2][0]]
            box_max_vector = [box[0][1], box[1][1], box[2][1]]
            box_vector = np.array(box_max_vector) - np.array(box_min_vector)
            box_num=np.floor(box_vector/R_search).astype(int) # make sure the box units is big rather than small. So that no point will be missed.
            box_len = box_vector/box_num
            N_atom=len(dirc_vec)
            loca_vec -= box_min_vector # the original program defines like this: all the location starts from zero point.
            loca_unit = [np.floor(vec / box_len).astype(int) % box_num for vec in loca_vec]
            atom_list = [Point_Vector(atom_index[i],loca_vec[i], dirc_vec[i], loca_unit[i] ) for i in range(N_atom)]

            ## Assign all loca_vec to units_box

            unit_box=create_3d_array(box_num[0],box_num[1],box_num[2])
            for atom in atom_list:
                x, y, z = atom.unit
                unit_box[x][y][z].append(atom)

            ## Iterate through all the atoms
            for atom0 in atom_list:
                # find all the unit_box id
                # the box id of the atom self
                S=0
                calculated_num=0
                
                pari_atom = adjacent_atom(atom0.unit, unit_box, box_num)

                for atom1 in pari_atom:
                    # 下面的一行是神，不论是谁看到这行代码，不要去试图搞懂它，请相信它，膜拜它，赞美它。
                    dist = np.linalg.norm(abs(atom0.loca - atom1.loca)-box_vector*atom1.PB_cross_flag)
                    if atom1.index != atom0.index and dist <= R_search:
                        calculated_num += 1
                        S += 1.5 * cos(atom0.vec, atom1.vec) ** 2 - 0.5

                if calculated_num > 0:
                    atom0.ord = S / calculated_num

            return atom_list

        def period_vec(point_before, point_after, box_xyz):
            delta = point_after - point_before
            box_lengths = np.array([x[1] - x[0] for x in box_xyz])
            # Vectorized boundary condition adjustment
            delta[delta > box_lengths / 2] -= box_lengths[delta > box_lengths / 2]
            delta[delta < -box_lengths / 2] += box_lengths[delta < -box_lengths / 2]
            return delta
        
        start_time = time.time()  # 记录开始时间
        
        # 检查前 10 个原子id是不是顺序排列的。
        for i in range(10):
            if self.atom_data[i+1][0]-self.atom_data[i][0]!=1:
                self.sort(0)
                
        # 检查是否有 chain_num 的输入
        if chain_num==0:
            # 检查是否有mol列
            if 'mol' in self.table_headers: #有mol列
                # 检查是否标记好了每条链
                all_elements_are_one = all(mol == self.atom_data['mol'][0] for mol in self.atom_data['mol'])
                if  all_elements_are_one:
                    # 没有则标记 而且没有chain_num输入
                    print('error: 无法从 mol 列推断 chain_num ,同时没有输入链数')
                    return False
                else:
                    chain_num=self.atom_num//sum(self.atom_data['mol']==self.atom_data['mol'][0])
                    self.mol_add(chain_num) #添加一个正确的mol 列
            else: # 无 mol 列
                # 没有则标记 而且没有chain_num输入
                print('error: 无 mol 列,同时没有输入链数')
                return False

        # check if ord list already exists
        if 'ord' in self.table_headers:
            self.delete('ord')

        # Compute the interval and the edge indices once
        N = self.atom_num // chain_num
        edge_head = np.arange(0, self.atom_num, N)
        edge_tail = edge_head + N - 1
        loca_vec = np.column_stack((self.atom_data['x'], self.atom_data['y'], self.atom_data['z']))
        dirc_vec = np.zeros_like(loca_vec)

        # Valid indices are all indices except the edges
        valid_indices = np.setdiff1d(np.arange(self.atom_num), np.concatenate([edge_head, edge_tail]))

        # Calculate direction vectors for valid indices
        for i in valid_indices:
            point_before = loca_vec[i - 1]
            point_after = loca_vec[i + 1]
            dirc_vec[i] = period_vec(point_before, point_after, self.box_xyz)

        # Handle edge cases by copying the direction vector from the nearest valid atom
        dirc_vec[edge_head] = dirc_vec[edge_head + 1]  # Assuming edge_head + 1 is a valid index
        dirc_vec[edge_tail] = dirc_vec[edge_tail - 1]  # Assuming edge_tail - 1 is a valid index

        atom_list = orderCalcu(self.atom_data['id'], loca_vec, dirc_vec, R_search, self.box_xyz)
        self.add('ord',[atom.ord for atom in atom_list])  

        end_time = time.time()  # 记录结束时间
        elapsed_time = end_time - start_time  # 计算经过的时间
        print(elapsed_time)

    def O_CB_add(self, R_search=5.4):
        
        class atom_unit:
            def __init__(self, index, loca, unit):
                self.index = index
                self.loca = loca
                self.unit = unit
                
        def O_CB_calcu(atom_index, loca_vec, R_search, box, l=4):
            from scipy.special import sph_harm
            def calculate_angles(atom0, atom1):
                r = atom1.loca - atom0.loca
                r_mag = np.linalg.norm(r)
                theta = np.arccos(r[2] / r_mag)  # z-axis is the third element
                phi = np.arctan2(r[1], r[0])  # y-axis is second, x-axis is first
                return theta, phi
            
            box_min_vector = [box[0][0], box[1][0], box[2][0]]
            box_max_vector = [box[0][1], box[1][1], box[2][1]]
            box_vector = np.array(box_max_vector) - np.array(box_min_vector)
            box_num=np.floor(box_vector/R_search).astype(int) # make sure the box units is big rather than small. So that no point will be missed.
            box_len = box_vector/box_num
            loca_vec -= box_min_vector # the original program defines like this: all the location starts from zero point.
            loca_unit = [np.floor(vec / box_len).astype(int) % box_num for vec in loca_vec]
            atom_list = [atom_unit(atom_index[i],loca_vec[i], loca_unit[i] ) for i in range(len(loca_vec))]

            # Assign all loca_vec to units_box
            unit_box = create_3d_array(box_num[0], box_num[1], box_num[2])
            for atom in atom_list:
                x, y, z = atom.unit
                unit_box[x][y][z].append(atom)

            # Calculation of the O_CB order parameter
            O_CB_values = []
            for atom0 in atom_list:
                atom_pari_list = adjacent_atom(atom0.unit, unit_box, box_num)
                
                Ql = []
                for atom_pari in atom_pari_list:
                    # 下面的一行是神，不论是谁看到这行代码，不要去试图搞懂它，请相信它，膜拜它，赞美它。
                    dist = np.linalg.norm(abs(atom0.loca - atom_pari.loca)-box_vector*atom_pari.PB_cross_flag)
                    if atom_pari.index != atom0.index and dist <= R_search:
                        theta, phi = calculate_angles(atom0, atom_pari)
                        for m in range(0, l+1):  # Adjusted range to cover negative m values
                            Ylm = sph_harm(m, l, phi, theta)
                            Ql.append(Ylm * np.conj(Ylm))
                Nb = len(atom_pari_list)
                if Nb > 0:
                    O_CB = (1 / Nb) * np.sum([(2 * np.pi * Ql_j / (l + 1))**0.5  for Ql_j in Ql])
                    O_CB_values.append(O_CB.real)  # Store the real part
                else:
                    O_CB_values.append(0)  # Handle cases with no neighbors

            return O_CB_values  # Return the average O_CB across all atoms

        start_time = time.time()  # 记录开始时间
        # check if ord list already exists
        if 'O_CB' in self.table_headers:
            self.delete('O_CB')
            
        loca_vec = np.column_stack((self.atom_data['x'], self.atom_data['y'], self.atom_data['z']))
        O_CB_list = O_CB_calcu(self.atom_data['id'], loca_vec, R_search, self.box_xyz)
        self.add('O_CB',O_CB_list)  

        end_time = time.time()  # 记录结束时间
        elapsed_time = end_time - start_time  # 计算经过的时间
        print(elapsed_time)

## 全局的两个函数中使用

def adjacent_atom(atom0_unit, unit_box, box_num):
    pari_atom = []
    cross_flag = [False, False, False]
    for i in [-1, 0, 1]:
        id_x = (atom0_unit[0] + i)
        cross_flag[0] = id_x in [-1, box_num[0]]
        for j in [-1, 0, 1]:
            id_y = (atom0_unit[1] + j)
            cross_flag[1] = id_y in [-1, box_num[1]]
            for k in [-1, 0, 1]:
                id_z = (atom0_unit[2] + k)
                cross_flag[2] = id_z in [-1, box_num[2]]
                sub_pari_atoms = unit_box[id_x % box_num[0]][id_y % box_num[1]][id_z % box_num[2]]
                for atom in sub_pari_atoms:
                    atom.PB_cross_flag = cross_flag[:]
                pari_atom += sub_pari_atoms
    return pari_atom

def create_3d_array(a, b, c):
        return [[[[ ] for _ in range(c)] for _ in range(b)] for _ in range(a)]
    

## MD 的 IO 部分

def MD_trajec_read(file_path, split_num=None, sort_column=None ):
    # Open and read all lines from the file
    with open(file_path, 'r') as file:
        lines = file.readlines()

    # Extract the number of atoms from the fourth line
    atom_num = int(lines[3].strip())

    # If split_num is specified, select lines at equal intervals
    if split_num :
        group_size = atom_num + 9
        total_groups = len(lines) // group_size  # Total number of complete groups in the file
        selected_indices = [i * group_size for i in range(total_groups)]  # Start index of each group

        # Now calculate the interval for selecting groups
        interval = max(1, total_groups // split_num)  # Avoid division by zero or negative values

        # Use the interval to select group start indices
        selected_group_starts = selected_indices[::interval][:split_num]

        # Collect the selected groups
        lines_temp = []
        for start in selected_group_starts:
            lines_temp.extend(lines[start:start + group_size])
        lines = lines_temp
    
    MD_trajectory = []
    blocks = re.split(r'ITEM: TIMESTEP\n', ''.join(lines))[1:]  # Split data into blocks

    for block in blocks:
        # The head format:x
            # ITEM: TIMESTEP
            # 15000000
            # ITEM: NUMBER OF ATOMS
            # 12000
            # ITEM: BOX BOUNDS pp pp pp
            # -2.2179263477672642e+00 6.2217926417155326e+01
            # -2.2179263477672642e+00 6.2217926417155326e+01
            # -2.2179263477672642e+00 6.2217926417155326e+01
            # ITEM: ATOMS id type x y z c_chains 

        lines = block.strip().split('\n')
        header_lines= ['ITEM: TIMESTEP']+lines[:7]
        timestep = int(lines[0].strip())  # Time or temperature
        num_atoms = int(lines[2].split()[-1])  # Number of atoms
        box_bounds = [tuple(map(float, line.split())) for line in lines[4:7]]  # Box dimensions
        
        # Initialize the dictionary for atom properties
        header_line_index = 7
        table_headers = lines[header_line_index].split()[2:]

        # Dynamically determine which columns should be treated as integers
        sample_data_line = lines[header_line_index + 1].strip().split()
        int_column_list = [i for i, value in enumerate(sample_data_line) if '.' not in value]

        # Define the data types for each column, defaulting to float unless the column index is in int_column_list
        dtypes = [(table_headers[i], 'int' if i in int_column_list else 'float') for i in range(len(table_headers))]
        atom_data = np.zeros(num_atoms, dtype=dtypes)  # Create structured array with specified data types

        # Fill in the data from the file into the structured array
        for i, line in enumerate(lines[header_line_index + 1:]):
            values = [float(i) for i in line.strip().split()]
            for j, value in enumerate(values):
                # Fill the appropriate column with the right type of value
                if j in int_column_list:
                    atom_data[table_headers[j]][i] = int(value)
                else:
                    atom_data[table_headers[j]][i] = float(value)

        # Calculate the box volume (assuming orthogonal bounds for simplicity)
        volume = 1
        for i in range(3):
            volume *= (box_bounds[i][1] - box_bounds[i][0])

        # Create MD_moment object
        md_moment = MD_moment(timestep, box_bounds, volume, num_atoms, table_headers, header_lines, atom_data)
        if sort_column is not None:
            md_moment.sort(sort_column)

        MD_trajectory.append(md_moment)

    return MD_trajectory

def MD_trajec_write(outfile_path, MD_trajectory):
    with open(outfile_path, 'w') as file:
        for md_moment in MD_trajectory:
            # Write the header lines
            for line in md_moment.header_lines:
                file.write(line+'\n')

            # Write the ITEM: ATOMS line followed by the table headers
            file.write("ITEM: ATOMS " + " ".join(md_moment.table_headers) + '\n')
            
            # Write the atom data
            for data in md_moment.atom_data:
                file.write(' '.join(map(str, data)) + '\n')

def MD_trajec_rho_enta2moment(trajectory, Z1_SP_path):
    ## 把Z1 code 的 Z1_SP 文件写到moment 里面
    import re
    # 函数来检测一行是否含有三个浮点数
    def is_coordinate_line(line):
        # 使用正则表达式匹配三个浮点数
        return bool(re.match(r'^\s*(-?\d+\.\d+\s+){2}-?\d+\.\d+\s*$', line))
    def is_number_line(line):
        # This function checks if the line contains only one number.
        # Adjust the logic based on the actual format of the number and line.
        try:
            float(line.strip())
            return True
        except ValueError:
            return False
        
    with open(Z1_SP_path, 'r') as file:
        lines = file.readlines()

    # 找出所有坐标行 cor_num ，坐标行的特征是一行有3个浮点数。以 cor_num-1 将lines分为若干block，
    # 所有bolck 放到 Z1_moment_list 里面。
    Z1_moment_list = []
    start_index = 0

    for i, line in enumerate(lines[1:]):
        if is_coordinate_line(line):
            # 如果是坐标行，将前一个块（从start_index到当前行的前一行）添加到Z1_moment_list
            if i > start_index:
                Z1_moment_list.append(lines[start_index+1:i])
            # 更新当前块的起始索引为当前行的下一行
            start_index = i + 1
    # 添加最后一个块（如果存在）
    if start_index < len(lines):
        Z1_moment_list.append(lines[start_index+1:])

    # 以每行包含一个数字的 列为界限，将 Z1_moment 分为 polymer_enta_point_list
    polymer_entaPoint_all_moment = []
    for j, block in enumerate(Z1_moment_list):
        polymer_entaPoint=[]
        temp_block = []
        for line in block:
            if is_number_line(line):
                if temp_block:
                    polymer_entaPoint.append(temp_block)
                    temp_block = []
            else:
                temp_block.append(line)
        if temp_block:  # Append the last sub-block if any
            polymer_entaPoint.append(temp_block)
        polymer_entaPoint_all_moment.append(polymer_entaPoint)

    moment_num=len(polymer_entaPoint_all_moment)
    del Z1_moment_list

    enta_point_all_moment=[]
    for one_moment in polymer_entaPoint_all_moment:
        # 遍历 polymer_enta_point_list 中的每一行
        polymer_entaPoint_one_moment=[]
        for polymer in one_moment:
            
            enta_point_list_one=[]
            # 从 enta_point_line 中读取倒数第二个元素并添加到 temp_list
            
            for enta_point in polymer:
                line_list=enta_point.strip().split()
                enta_point_list_one.append(int(float(line_list[-2])))
            polymer_entaPoint_one_moment.append(enta_point_list_one)
        enta_point_all_moment.append(polymer_entaPoint_one_moment)
    del polymer_entaPoint_all_moment

    # enta_point_list to enta_length_list
    for mome_id, point_list_In_one_moment in enumerate(enta_point_all_moment) :
        enta_length_list=[]
        enta_point_flag_list=[]
        for point_list in point_list_In_one_moment:
            # 填写 enta_point_flag_list
            for point_index in range(1,point_list[-1]+1):
                if point_index in point_list:
                    enta_point_flag_list.append(1)
                else:
                    enta_point_flag_list.append(0)

            # 填写 enta_length_list 第一个原子也要填写
            enta_length_list.append(point_list[1]-point_list[0])
            for index in range(1,len(point_list)):                
                times=point_list[index]-point_list[index-1]
                for i in range(times):
                    enta_length_list.append(times)
        
        # calculate rho_enta
        N=trajectory[0].atom_num
        V=trajectory[0].V
        rho_enta_list = [N / V / enta_length for enta_length in enta_length_list]
        trajectory[mome_id].add('enta_length',enta_length_list)
        trajectory[mome_id].add('rho_enta',rho_enta_list)
        # trajectory[mome_id].add('enta_point',enta_point_flag_list)

## 并行的部分

from multiprocessing import Pool, cpu_count
from functools import partial

def process_moment(R_search, chain_num, moment):
    # 假设order_para是一个已经定义好的方法，用于计算和更新每个moment的参数
    moment.orderPara_add(R_search = R_search, chain_num = chain_num)
    return moment
    
def parallel_write_order_to_MD_trajec(trajectory, R_search=7, chain_num=0):
    num_cores = cpu_count()
    print(f"Number of available cores: {num_cores}")
    
    # 使用functools.partial来固定process_moment函数中的R_search参数
    partial_process_moment = partial(process_moment, R_search, chain_num)
    
    # 创建一个进程池，利用所有可用核心
    with Pool(num_cores) as pool:
        # 传递partial_process_moment函数给pool.map
        updated_trajectory = pool.map(partial_process_moment, trajectory)
    
    return updated_trajectory

## 不依赖于 MD 类的各种计算函数

def calculate_correlation(list1, list2):
    if len(list1) != len(list2):
        raise ValueError("The lists must have the same number of elements")

    mean_list1 = np.mean(list1)
    mean_list2 = np.mean(list2)

    deviation_list1 = [x - mean_list1 for x in list1]
    deviation_list2 = [y - mean_list2 for y in list2]

    numerator = np.sum(np.multiply(deviation_list1, deviation_list2))
    denominator = np.sqrt(np.sum(np.square(deviation_list1)) * np.sum(np.square(deviation_list2)))

    return numerator / denominator if denominator != 0 else 0